"""
使用 Pinecone 实现基于向量数据库的 KNN 手写数字识别

依赖：
- pinecone-client
- scikit-learn
- numpy
- matplotlib
- pillow

功能：
1. 将 sklearn digits 数据集上传到 Pinecone
2. 使用 Pinecone 进行 KNN 查询
3. 显示查询结果（id, distance, label）
4. 预测数字并可视化
"""

from __future__ import annotations

import os
import time
from typing import List, Tuple, Optional
import numpy as np
from sklearn import datasets
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from PIL import Image

try:
    from pinecone import Pinecone, ServerlessSpec
    PINECONE_AVAILABLE = True
except ImportError:
    PINECONE_AVAILABLE = False
    print("警告: pinecone 未安装。请运行: pip install pinecone-client")


class PineconeKNN:
    """使用 Pinecone 实现的 KNN 分类器"""
    
    def __init__(self, api_key: Optional[str] = None, index_name: str = "digits-knn"):
        """
        初始化 Pinecone KNN 分类器
        
        Args:
            api_key: Pinecone API key (如果为 None，从环境变量 PINECONE_API_KEY 读取)
            index_name: Pinecone 索引名称
        """
        if not PINECONE_AVAILABLE:
            raise ImportError("请先安装 pinecone: pip install pinecone-client")
        
        self.api_key = api_key or os.getenv("PINECONE_API_KEY")
        if not self.api_key:
            raise ValueError(
                "请设置 Pinecone API key:\n"
                "1. 通过环境变量: export PINECONE_API_KEY='your-api-key'\n"
                "2. 或在代码中传入: PineconeKNN(api_key='your-api-key')"
            )
        
        self.index_name = index_name
        self.pc = Pinecone(api_key=self.api_key)
        self.index = None
        self.dimension = 64  # sklearn digits 是 8x8=64 维
        
    def create_index(self, dimension: int = 64, metric: str = "euclidean"):
        """
        创建 Pinecone 索引
        
        Args:
            dimension: 向量维度
            metric: 距离度量 (euclidean, cosine, dotproduct)
        """
        self.dimension = dimension
        
        # 检查索引是否已存在
        existing_indexes = [idx.name for idx in self.pc.list_indexes()]
        
        if self.index_name in existing_indexes:
            print(f"索引 '{self.index_name}' 已存在，将删除并重新创建...")
            self.pc.delete_index(self.index_name)
            time.sleep(1)  # 等待删除完成
        
        print(f"创建索引 '{self.index_name}'...")
        self.pc.create_index(
            name=self.index_name,
            dimension=dimension,
            metric=metric,
            spec=ServerlessSpec(
                cloud='aws',
                region='us-east-1'
            )
        )
        
        # 等待索引准备就绪
        while not self.pc.describe_index(self.index_name).status['ready']:
            time.sleep(1)
        
        self.index = self.pc.Index(self.index_name)
        print(f"索引创建成功！")
    
    def connect_to_index(self):
        """连接到已存在的索引"""
        existing_indexes = [idx.name for idx in self.pc.list_indexes()]
        
        if self.index_name not in existing_indexes:
            raise ValueError(f"索引 '{self.index_name}' 不存在，请先调用 create_index()")
        
        self.index = self.pc.Index(self.index_name)
        print(f"已连接到索引 '{self.index_name}'")
    
    def upload_data(self, X: np.ndarray, y: np.ndarray, batch_size: int = 100):
        """
        上传数据到 Pinecone
        
        Args:
            X: 特征向量 (n_samples, n_features)
            y: 标签 (n_samples,)
            batch_size: 批量上传大小
        """
        if self.index is None:
            raise ValueError("请先创建或连接到索引")
        
        n_samples = X.shape[0]
        print(f"开始上传 {n_samples} 个向量到 Pinecone...")
        
        # 批量上传
        for i in range(0, n_samples, batch_size):
            batch_end = min(i + batch_size, n_samples)
            batch_vectors = []
            
            for j in range(i, batch_end):
                vector_id = str(j)
                vector = X[j].tolist()
                label = int(y[j])
                
                batch_vectors.append({
                    "id": vector_id,
                    "values": vector,
                    "metadata": {"label": label}
                })
            
            self.index.upsert(vectors=batch_vectors)
            
            if (i + batch_size) % 500 == 0 or batch_end == n_samples:
                print(f"已上传 {batch_end}/{n_samples} 个向量")
        
        print("上传完成！")
        
        # 等待索引更新
        time.sleep(2)
        stats = self.index.describe_index_stats()
        print(f"索引统计: {stats}")
    
    def query(self, query_vector: np.ndarray, k: int = 10) -> Tuple[List[dict], int]:
        """
        查询最近邻并预测标签
        
        Args:
            query_vector: 查询向量 (1D array)
            k: 返回的最近邻数量
            
        Returns:
            matches: 匹配结果列表
            predicted_label: 预测的标签
        """
        if self.index is None:
            raise ValueError("请先创建或连接到索引")
        
        # 查询
        results = self.index.query(
            vector=query_vector.tolist(),
            top_k=k,
            include_metadata=True
        )
        
        matches = results['matches']
        
        # 投票决定预测标签
        labels = [match['metadata']['label'] for match in matches]
        predicted_label = int(np.bincount(labels).argmax())
        
        return matches, predicted_label
    
    def delete_index(self):
        """删除索引"""
        if self.index_name in [idx.name for idx in self.pc.list_indexes()]:
            self.pc.delete_index(self.index_name)
            print(f"索引 '{self.index_name}' 已删除")


def visualize_digit(digit_vector: np.ndarray, title: str = "Digit"):
    """
    可视化 8x8 数字图像（简洁清晰版，匹配图片效果）
    
    Args:
        digit_vector: 64维向量
        title: 图像标题
    """
    digit_image = digit_vector.reshape(8, 8)
    
    # 创建大窗口，清晰显示
    plt.figure(figsize=(6, 6), facecolor='white')
    
    # 使用 binary 色图和 nearest 插值，产生清晰的黑白块状效果
    plt.imshow(digit_image, cmap='binary', interpolation='nearest')
    
    # 简洁风格：无坐标轴
    plt.axis('off')
    
    # 大标题
    plt.title(title, fontsize=20, fontweight='bold', pad=10)
    
    plt.tight_layout()
    plt.show()


def print_query_results(matches: List[dict], predicted_label: int):
    """
    打印查询结果
    
    Args:
        matches: Pinecone 查询结果
        predicted_label: 预测的标签
    """
    print("\n" + "="*50)
    print("查询结果 (最近邻):")
    print("="*50)
    
    for match in matches:
        vector_id = match['id']
        distance = match['score']  # Pinecone 返回的是相似度分数
        label = match['metadata']['label']
        print(f"id: {vector_id:>4}, distance: {distance:>8.1f}, label: {label}.0")
    
    print("="*50)
    print(f"Predicted digit: {predicted_label}.0")
    print("="*50 + "\n")


def main():
    """主函数：演示 Pinecone KNN 的完整流程"""
    
    print("="*60)
    print("Pinecone KNN 手写数字识别示例")
    print("="*60 + "\n")
    
    # 1. 加载数据
    print("步骤 1: 加载 sklearn digits 数据集...")
    digits = datasets.load_digits()
    X, y = digits.data, digits.target
    print(f"数据集大小: {X.shape[0]} 个样本, 每个样本 {X.shape[1]} 维")
    print(f"标签范围: {y.min()} - {y.max()}\n")
    
    # 2. 划分数据集
    print("步骤 2: 划分训练集和测试集...")
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    print(f"训练集: {X_train.shape[0]} 个样本")
    print(f"测试集: {X_test.shape[0]} 个样本\n")
    
    # 3. 初始化 Pinecone
    print("步骤 3: 初始化 Pinecone...")
    print("提示: 请确保已设置环境变量 PINECONE_API_KEY")
    print("      或在代码中传入 api_key 参数\n")
    
    try:
        knn = PineconeKNN(index_name="digits-knn-demo")
    except ValueError as e:
        print(f"错误: {e}")
        print("\n如何获取 Pinecone API Key:")
        print("1. 访问 https://www.pinecone.io/")
        print("2. 注册/登录账号")
        print("3. 在控制台获取 API Key")
        print("4. 设置环境变量: export PINECONE_API_KEY='your-api-key'")
        return
    
    # 4. 创建索引并上传数据
    print("步骤 4: 创建 Pinecone 索引...")
    knn.create_index(dimension=64, metric="euclidean")
    
    print("\n步骤 5: 上传训练数据到 Pinecone...")
    knn.upload_data(X_train, y_train)
    
    # 5. 查询测试样本
    print("\n步骤 6: 查询测试样本...")
    test_idx = 100  # 选择一个测试样本
    query_vector = X_test[test_idx]
    true_label = y_test[test_idx]
    
    print(f"查询样本索引: {test_idx}")
    print(f"真实标签: {true_label}")
    
    # 执行查询
    k = 10  # 查询 10 个最近邻
    matches, predicted_label = knn.query(query_vector, k=k)
    
    # 打印结果
    print_query_results(matches, predicted_label)
    
    # 可视化
    print("步骤 7: 可视化查询样本...")
    visualize_digit(query_vector, title=f"Predicted digit: {predicted_label}.0")
    
    # 6. 评估准确率（可选）
    print("\n步骤 8: 评估模型准确率（在部分测试集上）...")
    n_test_samples = min(50, len(X_test))  # 测试前50个样本
    correct = 0
    
    for i in range(n_test_samples):
        _, pred = knn.query(X_test[i], k=k)
        if pred == y_test[i]:
            correct += 1
    
    accuracy = correct / n_test_samples
    print(f"准确率: {accuracy:.4f} ({correct}/{n_test_samples})")
    
    # 7. 清理（可选）
    print("\n是否删除索引? (y/n): ", end="")
    try:
        response = input().strip().lower()
        if response == 'y':
            knn.delete_index()
    except:
        print("\n保留索引")
    
    print("\n" + "="*60)
    print("示例完成！")
    print("="*60)


if __name__ == "__main__":
    main()
